# -*- coding: utf-8 -*-
#
# Copyright 2017 zhangxp, st
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import defaultdict
import os
from uuid import uuid1
import pandas as pd
from pymysql import Timestamp

from rqalpha.environment import Environment
from rqalpha.utils.arg_checker import apply_rules, verify_that
from rqalpha.core.events import EVENT

from rqalpha.api import export_as_api
import time
import json

from sqlalchemy import DATE, DATETIME, VARCHAR, create_engine


NEXTT_MYSQL_DB_HOST = os.getenv('NEXTT_MYSQL_DB_HOST', "192.168.12.110")

class DB_CONF:
    host = NEXTT_MYSQL_DB_HOST
    port = '3309'
    user = "admin"
    password = "admin"
    database = "nextt"

class PlotStore(object):

    def __init__(self, to_sql=False, conn_dsn=None):
        #
        self._inst = None
        self._plots = defaultdict(dict)
        self._klines = []
        #
        if conn_dsn is None:
            self._conn_dsn = 'mysql+pymysql://' + DB_CONF.user + ':' + DB_CONF.password + '@' + DB_CONF.host + ':' + DB_CONF.port + '/' + DB_CONF.database + '?charset=utf8'
        else:
            self._conn_dsn = conn_dsn
        self._env = Environment.get_instance()
        if self._env is not None:
            if to_sql:
                self._env.event_bus.add_listener(EVENT.POST_STRATEGY_RUN, self._dump_sql)
            else:
                self._env.event_bus.add_listener(EVENT.POST_STRATEGY_RUN, self._save_file)
        
    def init(self, inst):
        #
        self._inst = inst

    def add_plot(self, dt, plot_type, value):
        if dt not in self._plots:
            self._plots[dt] = {}
        content = self._plots[dt]
        if plot_type == 'ORDER':
            tmp = content[plot_type] if plot_type in content else {}
            tmp.update(value)
            pass
        elif plot_type == 'LINE':
            tmp = content[plot_type] if plot_type in content else {}
            tmp.update(value)
            pass
        elif plot_type == 'TEXT':
            tmp = content[plot_type] if plot_type in content else []
            tmp.append(value)
            pass
        elif plot_type == 'ICON':
            tmp = content[plot_type] if plot_type in content else []
            tmp.append(value)
            pass
        content[plot_type] = tmp


    def add_kline(self, trading_dt, open, high, low, close):
        timestamp = int(time.mktime(trading_dt.timetuple()))
        self._klines.append([timestamp, open, high, low, close, None, None, None, None])
        pass

    def get_klines(self):
        return self._klines

    def get_plots(self):
        return self._plots

    def _save_file(self, event):
        freq = self._env.config.base.frequency
        if 'strategy_name' not in self._env.config.base:
            strategy_name = self._env.config.base.strategy_file
        else:
            strategy_name = self._env.config.base.strategy_name
        # 信号文件
        file_name = "%s_%s_%s.txt" % (strategy_name.replace('.py', ''), self._inst, freq)
        with open(file_name, 'w') as f:
            for k, v in self._plots.items():
                f.write('%s %s\n' % (k, v))
        # K线文件
        file_name = "klines_%s_%s.txt" % (self._inst, freq)
        import pickle
        # 通过写入二进制模式打开文件
        with open(file_name, 'w') as f:
            json.dump(self._klines, f)
        pass


    def _dump_sql(self, event):
        #
        mysql_connect = create_engine(self._conn_dsn)
        #
        if 'strategy_name' not in self._env.config.base:
            strategy_name = os.path.basename(self._env.config.base.strategy_file).split(".")[0]
        else:
            strategy_name = self._env.config.base.strategy_name
        start_date = self._env.config.base.start_date
        end_date = self._env.config.base.end_date
        freq = self._env.config.base.frequency
        uuid = uuid1()
        plot_info = {'strategy_name':strategy_name, 'universe':self._inst, 'start_date':start_date, 'end_date':end_date, 'freq':freq, 'uuid':uuid}
        df = pd.DataFrame.from_dict([plot_info])
        df.to_sql('tb_plot_info', mysql_connect, if_exists='append', index=False, dtype={'strategy_name': VARCHAR(100), "start_date": DATE, "end_date": DATE, 'freq': VARCHAR(16)})
        # 信号数据
        plot_list = []
        for k, v in self._plots.items():
            plot_list.append({'datetime':k, \
                'ORDER':json.dumps(v['ORDER']) if 'ORDER' in v else None,
                'LINE':json.dumps(v['LINE']) if 'LINE' in v else None, 
                'TEXT':json.dumps(v['TEXT']) if 'TEXT' in v else None, 
                'ICON':json.dumps(v['ICON']) if 'ICON' in v else None})
        df = pd.DataFrame.from_records(plot_list)
        df['uuid'] = uuid
        df.to_sql('tb_plot_data', mysql_connect, if_exists='append', index=False, dtype={"datetime": DATETIME, 'uuid': VARCHAR(100)})
        mysql_connect.dispose()
        pass


if __name__ == '__main__':
    pass